"""
http://baostock.com/baostock/index.php/Python_API文档

create table st_cn_stock_trade_inc_d (
    id bigint unsigned NOT NULL AUTO_INCREMENT COMMENT '自增主键',
    dt date not null comment '日期',
    stock_code varchar(20) not null default '' comment '股票代码(ticker)',
    stock_name varchar(30) not null default '' comment '股票中文名',
    open_price double comment '开盘价',
    close_price double comment '收盘价',
    high_price double comment '盘中最高价',
    low_price double comment '盘中最低价',
    pre_close_price double comment '昨日收盘价',
    amount bigint comment '成交额(元)',
    volume bigint comment '成交量(股)',
    turnover_rate double comment '换手率',
    trade_status int comment '交易状态 {1:正常交易, 0:停牌}',
    is_st int comment '是否st {1:是, 0:否}',
    PRIMARY KEY (id),
    KEY index_dt (dt, stock_code)
) ENGINE=InnoDB COMMENT='A股股票交易日表'

select dt, count(1) from st_cn_stock_trade_inc_d group by dt order by dt;
"""

from common import Args, StockApi, Database

args = Args.parse()
api = StockApi()
db = Database()

dt = args.date
start_date = args.start_date
end_date = args.end_date

fields_map = {
    'date': 'dt',
    'code': 'stock_code',
    'code_name': 'stock_name',
    'open': 'open_price',
    'high': 'high_price',
    'low': 'low_price',
    'close': 'close_price',
    'preclose': 'pre_close_price',
    'volume': 'volume',
    'amount': 'amount',
    'turn': 'turnover_rate',
    'tradestatus': 'trade_status',
    'isST': 'is_st'
}

# 获取A股代码
source_sql = '''
select 
    stock_name, stock_code
from st_cn_stock_basic_d
where dt=%s
and security_type=1 -- 股票类型
'''

# 判断是否已经有数据
# 原则上，已经有的数据不会是错的
exist_sql = '''
select count(1) cnt from st_cn_stock_trade_inc_d
where dt=%s
and stock_code=%s
'''


for stock in db.execute(source_sql, dt):
    stock_name = stock['stock_name']
    stock_code = stock['stock_code']
    print('processing:', stock)
    result = api.query_history_k_data_plus(
        stock_code,
        'date,code,open,close,high,low,preclose,amount,volume,turn,tradestatus,isST',
        start_date=start_date, 
        end_date=end_date,
        frequency='d',
        adjustflag='3',
        fields_map=fields_map
    )
    for row in result:
        dt = row['dt']
        stock_code = row['stock_code']
        cnt = db.execute(exist_sql, (dt, stock_code))[0]['cnt']
        if cnt > 0:
            print(f'{stock_code} {dt} is ready done.')
            continue

        row['stock_name'] = stock_name
        db.execute('delete from st_cn_stock_trade_inc_d where dt=%s and stock_code=%s', (dt, stock_code))
        db.write('st_cn_stock_trade_inc_d', [row])
